#!/usr/bin/env python

# Copyright 2016 Google, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""This application demonstrates how to perform basic operations on blobs
(objects) in a Google Cloud Storage bucket.

For more information, see the README.md under /storage  and the documentation
at https://cloud.google.com/storage/docs.
"""

import argparse
import datetime
import pprint

# [START storage_upload_file]
from google.cloud import storage

# [END storage_upload_file]


def create_bucket(bucket_name):
    """Creates a new bucket."""
    storage_client = storage.Client()
    bucket = storage_client.create_bucket(bucket_name)
    print('Bucket {} created'.format(bucket.name))


def delete_bucket(bucket_name):
    """Deletes a bucket. The bucket must be empty."""
    storage_client = storage.Client()
    bucket = storage_client.get_bucket(bucket_name)
    bucket.delete()
    print('Bucket {} deleted'.format(bucket.name))


def enable_default_kms_key(bucket_name, kms_key_name):
    # [START storage_set_bucket_default_kms_key]
    """Sets a bucket's default KMS key."""
    storage_client = storage.Client()
    bucket = storage_client.get_bucket(bucket_name)
    bucket.default_kms_key_name = kms_key_name
    bucket.patch()

    print('Set default KMS key for bucket {} to {}.'.format(
        bucket.name,
        bucket.default_kms_key_name))
    # [END storage_set_bucket_default_kms_key]


def get_bucket_labels(bucket_name):
    """Prints out a bucket's labels."""
    storage_client = storage.Client()
    bucket = storage_client.get_bucket(bucket_name)
    labels = bucket.labels
    pprint.pprint(labels)


def add_bucket_label(bucket_name):
    """Add a label to a bucket."""
    storage_client = storage.Client()
    bucket = storage_client.get_bucket(bucket_name)

    labels = bucket.labels
    labels['example'] = 'label'
    bucket.labels = labels
    bucket.patch()

    print('Updated labels on {}.'.format(bucket.name))
    pprint.pprint(bucket.labels)


def remove_bucket_label(bucket_name):
    """Remove a label from a bucket."""
    storage_client = storage.Client()
    bucket = storage_client.get_bucket(bucket_name)

    labels = bucket.labels

    if 'example' in labels:
        del labels['example']

    bucket.labels = labels
    bucket.patch()

    print('Removed labels on {}.'.format(bucket.name))
    pprint.pprint(bucket.labels)


# [START storage_list_buckets]
def list_buckets():
    """Lists all buckets."""
    storage_client = storage.Client()
    buckets = storage_client.list_buckets()

    for bucket in buckets:
        print(bucket.name)
# [END storage_list_buckets]


def list_blobs(bucket_name):
    """Lists all the blobs in the bucket."""
    storage_client = storage.Client()

    # Note: Client.list_blobs requires at least package version 1.17.0.
    blobs = storage_client.list_blobs(bucket_name)

    for blob in blobs:
        print(blob.name)


def list_blobs_with_prefix(bucket_name, prefix, delimiter=None):
    """Lists all the blobs in the bucket that begin with the prefix.

    This can be used to list all blobs in a "folder", e.g. "public/".

    The delimiter argument can be used to restrict the results to only the
    "files" in the given "folder". Without the delimiter, the entire tree under
    the prefix is returned. For example, given these blobs:

        /a/1.txt
        /a/b/2.txt

    If you just specify prefix = '/a', you'll get back:

        /a/1.txt
        /a/b/2.txt

    However, if you specify prefix='/a' and delimiter='/', you'll get back:

        /a/1.txt

    """
    storage_client = storage.Client()

    # Note: Client.list_blobs requires at least package version 1.17.0.
    blobs = storage_client.list_blobs(bucket_name, prefix=prefix,
                                      delimiter=delimiter)

    print('Blobs:')
    for blob in blobs:
        print(blob.name)

    if delimiter:
        print('Prefixes:')
        for prefix in blobs.prefixes:
            print(prefix)


# [START storage_upload_file]
def upload_blob(bucket_name, source_file_name, destination_blob_name):
    """Uploads a file to the bucket."""
    storage_client = storage.Client()
    bucket = storage_client.get_bucket(bucket_name)
    blob = bucket.blob(destination_blob_name)

    blob.upload_from_filename(source_file_name)

    print('File {} uploaded to {}.'.format(
        source_file_name,
        destination_blob_name))
# [END storage_upload_file]


def upload_blob_with_kms(bucket_name, source_file_name, destination_blob_name,
                         kms_key_name):
    # [START storage_upload_with_kms_key]
    """Uploads a file to the bucket, encrypting it with the given KMS key."""
    storage_client = storage.Client()
    bucket = storage_client.get_bucket(bucket_name)
    blob = bucket.blob(destination_blob_name, kms_key_name=kms_key_name)
    blob.upload_from_filename(source_file_name)

    print('File {} uploaded to {} with encryption key {}.'.format(
        source_file_name,
        destination_blob_name,
        kms_key_name))
    # [END storage_upload_with_kms_key]


def download_blob(bucket_name, source_blob_name, destination_file_name):
    """Downloads a blob from the bucket."""
    storage_client = storage.Client()
    bucket = storage_client.get_bucket(bucket_name)
    blob = bucket.blob(source_blob_name)

    blob.download_to_filename(destination_file_name)

    print('Blob {} downloaded to {}.'.format(
        source_blob_name,
        destination_file_name))


def delete_blob(bucket_name, blob_name):
    """Deletes a blob from the bucket."""
    storage_client = storage.Client()
    bucket = storage_client.get_bucket(bucket_name)
    blob = bucket.blob(blob_name)

    blob.delete()

    print('Blob {} deleted.'.format(blob_name))


def blob_metadata(bucket_name, blob_name):
    """Prints out a blob's metadata."""
    storage_client = storage.Client()
    bucket = storage_client.get_bucket(bucket_name)
    blob = bucket.get_blob(blob_name)

    print('Blob: {}'.format(blob.name))
    print('Bucket: {}'.format(blob.bucket.name))
    print('Storage class: {}'.format(blob.storage_class))
    print('ID: {}'.format(blob.id))
    print('Size: {} bytes'.format(blob.size))
    print('Updated: {}'.format(blob.updated))
    print('Generation: {}'.format(blob.generation))
    print('Metageneration: {}'.format(blob.metageneration))
    print('Etag: {}'.format(blob.etag))
    print('Owner: {}'.format(blob.owner))
    print('Component count: {}'.format(blob.component_count))
    print('Crc32c: {}'.format(blob.crc32c))
    print('md5_hash: {}'.format(blob.md5_hash))
    print('Cache-control: {}'.format(blob.cache_control))
    print('Content-type: {}'.format(blob.content_type))
    print('Content-disposition: {}'.format(blob.content_disposition))
    print('Content-encoding: {}'.format(blob.content_encoding))
    print('Content-language: {}'.format(blob.content_language))
    print('Metadata: {}'.format(blob.metadata))
    print("Temporary hold: ",
          'enabled' if blob.temporary_hold else 'disabled')
    print("Event based hold: ",
          'enabled' if blob.event_based_hold else 'disabled')
    if blob.retention_expiration_time:
        print("retentionExpirationTime: {}"
              .format(blob.retention_expiration_time))


def bucket_metadata(bucket_name):
    """Prints out a bucket's metadata."""
    # [START storage_get_bucket_metadata]
    storage_client = storage.Client()
    bucket = storage_client.get_bucket(bucket_name)

    print('ID: {}'.format(bucket.id))
    print('Name: {}'.format(bucket.name))
    print('Storage Class: {}'.format(bucket.storage_class))
    print('Location: {}'.format(bucket.location))
    print('Location Type: {}'.format(bucket.location_type))
    print('Cors: {}'.format(bucket.cors))
    print('Default Event Based Hold: {}'
          .format(bucket.default_event_based_hold))
    print('Default KMS Key Name: {}'.format(bucket.default_kms_key_name))
    print('Metageneration: {}'.format(bucket.metageneration))
    print('Retention Effective Time: {}'
          .format(bucket.retention_policy_effective_time))
    print('Retention Period: {}'.format(bucket.retention_period))
    print('Retention Policy Locked: {}'.format(bucket.retention_policy_locked))
    print('Requester Pays: {}'.format(bucket.requester_pays))
    print('Self Link: {}'.format(bucket.self_link))
    print('Time Created: {}'.format(bucket.time_created))
    print('Versioning Enabled: {}'.format(bucket.versioning_enabled))
    print('Labels:')
    pprint.pprint(bucket.labels)
    # [END storage_get_bucket_metadata]


def make_blob_public(bucket_name, blob_name):
    """Makes a blob publicly accessible."""
    storage_client = storage.Client()
    bucket = storage_client.get_bucket(bucket_name)
    blob = bucket.blob(blob_name)

    blob.make_public()

    print('Blob {} is publicly accessible at {}'.format(
        blob.name, blob.public_url))


def generate_signed_url(bucket_name, blob_name):
    """Generates a v2 signed URL for downloading a blob.

    Note that this method requires a service account key file. You can not use
    this if you are using Application Default Credentials from Google Compute
    Engine or from the Google Cloud SDK.
    """
    storage_client = storage.Client()
    bucket = storage_client.get_bucket(bucket_name)
    blob = bucket.blob(blob_name)

    url = blob.generate_signed_url(
        # This URL is valid for 1 hour
        expiration=datetime.timedelta(hours=1),
        # Allow GET requests using this URL.
        method='GET')

    print('The signed url for {} is {}'.format(blob.name, url))
    return url


# [START storage_generate_signed_url_v4]
def generate_download_signed_url_v4(bucket_name, blob_name):
    """Generates a v4 signed URL for downloading a blob.

    Note that this method requires a service account key file. You can not use
    this if you are using Application Default Credentials from Google Compute
    Engine or from the Google Cloud SDK.
    """
    storage_client = storage.Client()
    bucket = storage_client.get_bucket(bucket_name)
    blob = bucket.blob(blob_name)

    url = blob.generate_signed_url(
        version='v4',
        # This URL is valid for 15 minutes
        expiration=datetime.timedelta(minutes=15),
        # Allow GET requests using this URL.
        method='GET')

    print('Generated GET signed URL:')
    print(url)
    print('You can use this URL with any user agent, for example:')
    print('curl \'{}\''.format(url))
    return url
# [END storage_generate_signed_url_v4]


# [START storage_generate_upload_signed_url_v4]
def generate_upload_signed_url_v4(bucket_name, blob_name):
    """Generates a v4 signed URL for uploading a blob using HTTP PUT.

    Note that this method requires a service account key file. You can not use
    this if you are using Application Default Credentials from Google Compute
    Engine or from the Google Cloud SDK.
    """
    storage_client = storage.Client()
    bucket = storage_client.get_bucket(bucket_name)
    blob = bucket.blob(blob_name)

    url = blob.generate_signed_url(
        version='v4',
        # This URL is valid for 15 minutes
        expiration=datetime.timedelta(minutes=15),
        # Allow GET requests using this URL.
        method='PUT',
        content_type='application/octet-stream')

    print('Generated PUT signed URL:')
    print(url)
    print('You can use this URL with any user agent, for example:')
    print("curl -X PUT -H 'Content-Type: application/octet-stream' "
          "--upload-file my-file '{}'".format(url))
    return url
# [END storage_generate_upload_signed_url_v4]


def rename_blob(bucket_name, blob_name, new_name):
    """Renames a blob."""
    storage_client = storage.Client()
    bucket = storage_client.get_bucket(bucket_name)
    blob = bucket.blob(blob_name)

    new_blob = bucket.rename_blob(blob, new_name)

    print('Blob {} has been renamed to {}'.format(
        blob.name, new_blob.name))


def copy_blob(bucket_name, blob_name, new_bucket_name, new_blob_name):
    """Copies a blob from one bucket to another with a new name."""
    storage_client = storage.Client()
    source_bucket = storage_client.get_bucket(bucket_name)
    source_blob = source_bucket.blob(blob_name)
    destination_bucket = storage_client.get_bucket(new_bucket_name)

    new_blob = source_bucket.copy_blob(
        source_blob, destination_bucket, new_blob_name)

    print('Blob {} in bucket {} copied to blob {} in bucket {}.'.format(
        source_blob.name, source_bucket.name, new_blob.name,
        destination_bucket.name))


def bucket_commands(args):
    if args.command == 'list-buckets':
        list_buckets()
    elif args.command == 'bucket-metadata':
        bucket_metadata(args.bucket_name)
    elif args.command == 'create-bucket':
        create_bucket(args.bucket_name)
    elif args.command == 'enable-default-kms-key':
        enable_default_kms_key(args.bucket_name, args.kms_key_name)
    elif args.command == 'delete-bucket':
        delete_bucket(args.bucket_name)
    elif args.command == 'get-bucket-labels':
        get_bucket_labels(args.bucket_name)
    elif args.command == 'add-bucket-label':
        add_bucket_label(args.bucket_name)
    elif args.command == 'remove-bucket-label':
        remove_bucket_label(args.bucket_name)


def blob_commands(args):
    if args.command == 'list':
        list_blobs(args.bucket_name)
    elif args.command == 'list-with-prefix':
        list_blobs_with_prefix(args.bucket_name, args.prefix, args.delimiter)
    elif args.command == 'upload':
        upload_blob(
            args.bucket_name,
            args.source_file_name,
            args.destination_blob_name)
    elif args.command == 'upload-with-kms-key':
        upload_blob_with_kms(
            args.bucket_name,
            args.source_file_name,
            args.destination_blob_name,
            args.kms_key_name)
    elif args.command == 'download':
        download_blob(
            args.bucket_name,
            args.source_blob_name,
            args.destination_file_name)
    elif args.command == 'delete':
        delete_blob(args.bucket_name, args.blob_name)
    elif args.command == 'metadata':
        blob_metadata(args.bucket_name, args.blob_name)
    elif args.command == 'make-public':
        make_blob_public(args.bucket_name, args.blob_name)
    elif args.command == 'signed-url':
        generate_signed_url(args.bucket_name, args.blob_name)
    elif args.command == 'signed-url-download-v4':
        generate_download_signed_url_v4(args.bucket_name, args.blob_name)
    elif args.command == 'signed-url-upload-v4':
        generate_upload_signed_url_v4(args.bucket_name, args.blob_name)
    elif args.command == 'rename':
        rename_blob(args.bucket_name, args.blob_name, args.new_name)
    elif args.command == 'copy':
        copy_blob(
            args.bucket_name,
            args.blob_name,
            args.new_bucket_name,
            args.new_blob_name)


def main():
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)

    subparsers = parser.add_subparsers(dest='command')

    subparsers.add_parser(
        'list-buckets', help=list_buckets.__doc__)

    create_bucket_parser = subparsers.add_parser(
        'create-bucket', help=create_bucket.__doc__)
    create_bucket_parser.add_argument(
        'bucket_name', help='Your cloud storage bucket.')

    delete_create_parser = subparsers.add_parser(
        'delete-bucket', help=delete_bucket.__doc__)
    delete_create_parser.add_argument(
        'bucket_name', help='Your cloud storage bucket.')

    get_bucket_labels_parser = subparsers.add_parser(
        'get-bucket-labels', help=get_bucket_labels.__doc__)
    get_bucket_labels_parser.add_argument(
        'bucket_name', help='Your cloud storage bucket.')

    add_bucket_label_parser = subparsers.add_parser(
        'add-bucket-label', help=add_bucket_label.__doc__)
    add_bucket_label_parser.add_argument(
        'bucket_name', help='Your cloud storage bucket.')

    remove_bucket_label_parser = subparsers.add_parser(
        'remove-bucket-label', help=remove_bucket_label.__doc__)
    remove_bucket_label_parser.add_argument(
        'bucket_name', help='Your cloud storage bucket.')

    list_blobs_parser = subparsers.add_parser(
        'list', help=list_blobs.__doc__)
    list_blobs_parser.add_argument(
        'bucket_name', help='Your cloud storage bucket.')

    bucket_metadata_parser = subparsers.add_parser(
        'bucket-metadata', help=bucket_metadata.__doc__)
    bucket_metadata_parser.add_argument(
        'bucket_name', help='Your cloud storage bucket.')

    list_with_prefix_parser = subparsers.add_parser(
        'list-with-prefix', help=list_blobs_with_prefix.__doc__)
    list_with_prefix_parser.add_argument(
        'bucket_name', help='Your cloud storage bucket.')
    list_with_prefix_parser.add_argument('prefix')
    list_with_prefix_parser.add_argument('--delimiter', default=None)

    upload_parser = subparsers.add_parser(
        'upload', help=upload_blob.__doc__)
    upload_parser.add_argument(
        'bucket_name', help='Your cloud storage bucket.')
    upload_parser.add_argument('source_file_name')
    upload_parser.add_argument('destination_blob_name')

    enable_default_kms_parser = subparsers.add_parser(
        'enable-default-kms-key', help=enable_default_kms_key.__doc__)
    enable_default_kms_parser.add_argument(
        'bucket_name', help='Your cloud storage bucket.')
    enable_default_kms_parser.add_argument('kms_key_name')

    upload_kms_parser = subparsers.add_parser(
        'upload-with-kms-key', help=upload_blob_with_kms.__doc__)
    upload_kms_parser.add_argument(
        'bucket_name', help='Your cloud storage bucket.')
    upload_kms_parser.add_argument('source_file_name')
    upload_kms_parser.add_argument('destination_blob_name')
    upload_kms_parser.add_argument('kms_key_name')

    download_parser = subparsers.add_parser(
        'download', help=download_blob.__doc__)
    download_parser.add_argument(
        'bucket_name', help='Your cloud storage bucket.')
    download_parser.add_argument('source_blob_name')
    download_parser.add_argument('destination_file_name')

    delete_parser = subparsers.add_parser(
        'delete', help=delete_blob.__doc__)
    delete_parser.add_argument(
        'bucket_name', help='Your cloud storage bucket.')
    delete_parser.add_argument('blob_name')

    metadata_parser = subparsers.add_parser(
        'metadata', help=blob_metadata.__doc__)
    metadata_parser.add_argument(
        'bucket_name', help='Your cloud storage bucket.')
    metadata_parser.add_argument('blob_name')

    make_public_parser = subparsers.add_parser(
        'make-public', help=make_blob_public.__doc__)
    make_public_parser.add_argument(
        'bucket_name', help='Your cloud storage bucket.')
    make_public_parser.add_argument('blob_name')

    signed_url_parser = subparsers.add_parser(
        'signed-url', help=generate_signed_url.__doc__)
    signed_url_parser.add_argument(
        'bucket_name', help='Your cloud storage bucket.')
    signed_url_parser.add_argument('blob_name')

    signed_url_download_v4_parser = subparsers.add_parser(
        'signed-url-download-v4', help=generate_download_signed_url_v4.__doc__)
    signed_url_download_v4_parser.add_argument('blob_name')

    signed_url_upload_v4_parser = subparsers.add_parser(
        'signed-url-upload-v4', help=generate_upload_signed_url_v4.__doc__)
    signed_url_upload_v4_parser.add_argument('blob_name')

    rename_parser = subparsers.add_parser(
        'rename', help=rename_blob.__doc__)
    rename_parser.add_argument(
        'bucket_name', help='Your cloud storage bucket.')
    rename_parser.add_argument('blob_name')
    rename_parser.add_argument('new_name')

    copy_parser = subparsers.add_parser('copy', help=rename_blob.__doc__)
    copy_parser.add_argument('bucket_name', help='Your cloud storage bucket.')
    copy_parser.add_argument('blob_name')
    copy_parser.add_argument('new_bucket_name')
    copy_parser.add_argument('new_blob_name')

    args = parser.parse_args()
    bucket_commands(args)
    blob_commands(args)


if __name__ == '__main__':
    main()
